import xml.etree.ElementTree as ElementTree
from lxml.etree import Element, SubElement, tostring


# ----------------------------------------------------------------------------
def read(xml_path):
    root = ElementTree.parse(xml_path)
    # read image info
    image_info = {}
    folder = root.find('folder').text
    filename = root.find('filename').text

    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)
    depth = int(size.find('depth').text)

    image_info['width'] = width
    image_info['height'] = height
    image_info['depth'] = depth
    image_info['folder'] = folder
    image_info['filename'] = filename

    # read objects 
    objs = []
    for obj in root.findall('object'):
        name = obj.find('name').text
        bndbox = obj.find('bndbox')
        xmin = int(float(bndbox.find('xmin').text)) - 1
        ymin = int(float(bndbox.find('ymin').text)) - 1
        xmax = int(float(bndbox.find('xmax').text)) - 1
        ymax = int(float(bndbox.find('ymax').text)) - 1
        if xmin < 0: xmin = 0
        if ymin < 0: ymin = 0
        if xmax >= width - 1: xmax = width - 1
        if ymax >= height - 1: ymax = height - 1
        objs.append([name, xmin, ymin, xmax, ymax])
    return image_info, objs
 
def write(xml_path, image_info, objs):
    node_root = Element('annotation')
    # write image info
    node_folder = SubElement(node_root, 'folder')
    node_folder.text = image_info['folder']
    node_filename = SubElement(node_root, 'filename')
    node_filename.text = image_info['filename']
    node_size = SubElement(node_root, 'size')
    node_width = SubElement(node_size, 'width')
    node_width.text = '%s' % image_info['width']
    node_height = SubElement(node_size, 'height')
    node_height.text = '%s' % image_info['height']
    node_depth = SubElement(node_size, 'depth')
    node_depth.text = '%s' % image_info['depth']

    # write objects
    for obj in annotations:
        node_obj = SubElement(node_root, 'object')
        node_name = SubElement(node_obj, 'name')
        node_name.text = obj[0]
        node_pose = SubElement(node_obj, 'pose')
        node_pose.text = 'None'
        node_truncated = SubElement(node_obj, 'truncated')
        node_truncated.text = '0'
        node_difficult = SubElement(node_obj, 'difficult')
        node_difficult.text = '0'
        node_bndbox = SubElement(node_obj, 'bndbox')
        node_xmin = SubElement(node_bndbox, 'xmin')
        node_xmin.text = '%s' % (obj[1] + 1)
        node_ymin = SubElement(node_bndbox, 'ymin')
        node_ymin.text = '%s' % (obj[2] + 1)
        node_xmax = SubElement(node_bndbox, 'xmax')
        node_xmax.text = '%s' % (obj[3] + 1)
        node_ymax = SubElement(node_bndbox, 'ymax')
        node_ymax.text = '%s' % (obj[4] + 1)

    xml = tostring(node_root, pretty_print=True)
    with open(annot_path, 'w') as f:
        f.write(xml.decode('utf-8'))



# --------------------------------------------------------------------------
def read_yolov3(text_path):
    data = []
    with open(text_path, 'r') as f:
        for line in f.readlines():
            vec = line.strip().split()
            image_path = vec[0]
            boxes = np.array([np.array(list(map(int, b.split(',')))) for b in vec[1:]])
            data.append([image_path, boxes])
    return data


def write_yolov3(data, text_path, class_names=None):
    with open(text_path, 'w') as f:
        for image_path, boxes in data:
            boxes = 




